/*
 * Copyright 2019 Amazon.com, Inc. or its affiliates. All Rights Reserved.
 *
 * Licensed under the Apache License, Version 2.0 (the "License").
 * You may not use this file except in compliance with the License.
 * A copy of the License is located at
 *
 *  http://aws.amazon.com/apache2.0
 *
 * or in the "license" file accompanying this file. This file is distributed
 * on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either
 * express or implied. See the License for the specific language governing
 * permissions and limitations under the License.
 */
using System;
using System.Collections.Generic;
using System.Linq;

namespace Amazon.Lambda.RuntimeSupport
{
    /// <summary>
    /// Headers that are specific to each invocation that are used in varies parts of RuntimeSupport like the logging system.
    /// </summary>
    public interface IRuntimeApiHeaders
    {
        /// <summary>
        /// The AWS request ID associated with the request.
        /// This is the same ID returned to the client that called invoke().
        /// This ID is reused for retries on the same request.
        /// </summary>
        string AwsRequestId { get; }

        /// <summary>
        /// The trace id generated by Lambda for distributed tracing across AWS services.
        /// </summary>
        string TraceId { get; }

        /// <summary>
        /// Gets the tenant id for the Lambda function.
        /// </summary>
        string TenantId { get; }
    }

    internal class RuntimeApiHeaders : IRuntimeApiHeaders
    {
        internal const string HeaderAwsRequestId = "Lambda-Runtime-Aws-Request-Id";
        internal const string HeaderTraceId = "Lambda-Runtime-Trace-Id";
        internal const string HeaderClientContext = "Lambda-Runtime-Client-Context";
        internal const string HeaderCognitoIdentity = "Lambda-Runtime-Cognito-Identity";
        internal const string HeaderDeadlineMs = "Lambda-Runtime-Deadline-Ms";
        internal const string HeaderInvokedFunctionArn = "Lambda-Runtime-Invoked-Function-Arn";
        internal const string HeaderAwsTenantId = "Lambda-Runtime-Aws-Tenant-Id";

        public RuntimeApiHeaders(Dictionary<string, IEnumerable<string>> headers)
        {
            var caseInsensitiveHeaders = new Dictionary<string, IEnumerable<string>>(headers, StringComparer.OrdinalIgnoreCase);

            DeadlineMs = GetHeaderValueOrNull(caseInsensitiveHeaders, HeaderDeadlineMs);
            AwsRequestId = GetHeaderValueRequired(caseInsensitiveHeaders, HeaderAwsRequestId);
            ClientContextJson = GetHeaderValueOrNull(caseInsensitiveHeaders, HeaderClientContext);
            CognitoIdentityJson = GetHeaderValueOrNull(caseInsensitiveHeaders, HeaderCognitoIdentity);
            InvokedFunctionArn = GetHeaderValueOrNull(caseInsensitiveHeaders, HeaderInvokedFunctionArn);
            TraceId = GetHeaderValueOrNull(caseInsensitiveHeaders, HeaderTraceId);
            TenantId = GetHeaderValueOrNull(caseInsensitiveHeaders, HeaderAwsTenantId);
        }

        public string AwsRequestId { get; private set; }
        public string InvokedFunctionArn { get; private set; }
        public string TraceId { get; private set; }
        public string ClientContextJson { get; private set; }
        public string CognitoIdentityJson { get; private set; }
        public string DeadlineMs { get; private set; }
        public string TenantId { get; private set; }

        private string GetHeaderValueRequired(Dictionary<string, IEnumerable<string>> headers, string header)
        {
            return headers[header].FirstOrDefault();
        }

        private string GetHeaderValueOrNull(Dictionary<string, IEnumerable<string>> headers, string header)
        {
            if (headers.TryGetValue(header, out var values))
            {
                return values.FirstOrDefault();
            }

            return null;
        }
    }
}
